package chatgpt.util;

import cn.hutool.core.util.StrUtil;
import com.gpt.article.service.impl.taskHandle.task.GenQContentTasker;
import com.knuddels.jtokkit.Encodings;
import com.knuddels.jtokkit.api.Encoding;
import com.knuddels.jtokkit.api.EncodingRegistry;
import chatgpt.entity.chat.ChatCompletion;
import chatgpt.entity.chat.Message;
import com.ruoyi.common.core.domain.entity.SysDictData;
import com.ruoyi.system.service.ISysDictTypeService;
import lombok.experimental.UtilityClass;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

@Service
@UtilityClass
public class TokensUtil {

    private static final Map<String, Encoding> modelEncodingMap = new HashMap<>();
    private static final EncodingRegistry encodingRegistry = Encodings.newDefaultEncodingRegistry();
    private static ISysDictTypeService sysDictTypeService;
    @Autowired
    public void setISysDictTypeService(ISysDictTypeService sysDictTypeService) {
        TokensUtil.sysDictTypeService = sysDictTypeService;
    }
    static {
        for (ChatCompletion.Model model : ChatCompletion.Model.values()) {
            Optional<Encoding> encodingForModel = encodingRegistry.getEncodingForModel(model.getName());
            encodingForModel.ifPresent(encoding -> modelEncodingMap.put(model.getName(), encoding));
        }
        //hzb+ 添加自定义模型
        List<SysDictData> dictModellist = sysDictTypeService.selectDictDataByType("gpt_model");
        loop1:for (SysDictData sysDictData : dictModellist) {
            //去除ChatCompletion.Model.values() 内的枚举
            Boolean isFind = false;
            String modelName = sysDictData.getDictValue();
            loop2:for (ChatCompletion.Model model : ChatCompletion.Model.values()) {
                if (model.getName().equals(modelName)) {
                    isFind = true;
                    break loop2;
                }
            }
            if(!isFind){
                Optional<Encoding> encodingForModel = encodingRegistry.getEncodingForModel(sysDictData.getDictValue());
                encodingForModel.ifPresent(encoding -> modelEncodingMap.put(sysDictData.getDictValue(), encoding));
            }

        }
    }

    /**
     * 计算tokens
     * @param modelName 模型名称
     * @param messages 消息列表
     * @return 计算出的tokens数量
     */

    public static int tokens(String modelName, List<Message> messages) {
        Encoding encoding = modelEncodingMap.get(modelName);
        if (encoding == null) {
            throw new IllegalArgumentException("Unsupported model: " + modelName);
        }

        int tokensPerMessage = 0;
        int tokensPerName = 0;
        if (modelName.startsWith("gpt-4")) {
            tokensPerMessage = 3;
            tokensPerName = 1;
        } else if (modelName.startsWith("gpt-3.5-turbo")) {
            tokensPerMessage = 4; // every message follows <|start|>{role/name}\n{content}<|end|>\n
            tokensPerName = -1; // if there's a name, the role is omitted
        }
        int sum = 0;
        for (Message message : messages) {
            sum += tokensPerMessage;
            sum += encoding.countTokens(message.getContent());
            sum += encoding.countTokens(message.getRole());
            if (StrUtil.isNotBlank(message.getName())) {
                sum += encoding.countTokens(message.getName());
                sum += tokensPerName;
            }
        }
        sum += 3;
        return sum;
    }
}